package com.lagou.rpc.client.netty;

import com.lagou.rpc.common.exceptions.SimpleRpcException;
import io.netty.channel.Channel;
import io.netty.channel.ChannelId;

import java.util.Map;
import java.util.WeakHashMap;

/**
 * 请求结果注册器
 *
 * @author wlz
 * @date 2020/5/1
 */
public class RequestResultRegistry {
    private Map<ChannelId, Object> requestResultMap = new WeakHashMap<>();


    public void registerResultAndNotify(Channel channel, Object result) {
        synchronized (channel) {
            ChannelId channelId = channel.id();
            requestResultMap.put(channelId, result);
            channel.notifyAll();
        }
    }

    public Object waitForResult(Channel channel, long timeout) {
        synchronized (channel) {
            ChannelId channelId = channel.id();
            // 首先判断结果是否已经返回
            if (requestResultMap.containsKey(channelId)) {
                // 如果已经返回，则直接返回结果
                return requestResultMap.remove(channelId);
            } else {
                try {
                    // 否则进行等待
                    channel.wait(timeout);
                } catch (InterruptedException e) {
                    throw new SimpleRpcException("wait for result failed", e);
                }
                // 等待结束后再尝试获取结果
                if (requestResultMap.containsKey(channelId)) {
                    // 如果有，则直接返回
                    return requestResultMap.remove(channelId);
                } else {
                    // 如果没有，则表示是等待超时
                    throw new SimpleRpcException("wait for result timeout");
                }
            }

        }
    }
}
